import numpy as np
import scipy.io.wavfile
import scipy.fftpack
import matplotlib.pyplot as plt
import seaborn as sns
import os
import sys


def sep(label = '', cnt=32):
    print('-' * cnt, label, '-' * cnt, sep='')


sep('Load')
BASE_DIR, FILE_NAME = os.path.split(__file__)
path = '../../../../../large_data/audio/zsn-stop1.wav'
AUDIO_PATH = os.path.join(BASE_DIR, path)
sr, signal = scipy.io.wavfile.read(AUDIO_PATH)
print('signal_len', signal.shape, signal.dtype)
if len(signal.shape) >= 2:
    signal = signal[:, 0]
print('signal_len', signal.shape, signal.dtype)

sep('Emphasize')
signal = np.append(signal[0:1], signal[1:] - signal[:-1] * 0.97)

sep('Frames splitting')
signal_len = len(signal)
print('signal_len', signal_len)
frame_len = int(np.round(sr / 40))
frame_stride = int(np.round(sr / 100))
n_frames = int(np.ceil(np.absolute(signal_len - frame_len) / frame_stride))
padded = n_frames * frame_stride + frame_len
signal = np.append(signal, np.zeros(padded - signal_len, dtype=signal.dtype))
print('signal_len_padded', signal_len)
print('frame_len', frame_len)
print('frame_stride', frame_stride)
print('n_frames', n_frames)
idx_col = np.arange(0, frame_len).reshape(1, -1)
idx_row = np.arange(0, n_frames * frame_stride, frame_stride).reshape(-1, 1)
idx = np.asarray(idx_row + idx_col, dtype=np.int64)
frames = signal[idx]
print('frames', frames.shape, frames.dtype)

sep('Add windows')
win = np.hamming(frame_len)
frames *= win

sep('Fourier Transform')
NFFT = 512
rfft = np.fft.rfft(frames, NFFT)
rfft = np.absolute(rfft)
print('rfft', rfft.shape, rfft.dtype)
_, n_spec = rfft.shape

sep('Power spectrum')
pow_frames = rfft ** 2 / NFFT

sep('Mel filters')


def hz2mel(hz):
    return 2595 * np.log10(1 + hz / 700)


def mel2hz(mel):
    return (10 ** (mel / 2595) - 1) * 700


n_filters = 40
low_hz = 0
high_hz = sr / 2
low_mel = 0
high_mel = hz2mel(high_hz)
mel_arr = np.linspace(low_mel, high_mel, n_filters + 2)
hz_arr = mel2hz(mel_arr)
y1y2 = np.tile([[100], [-100]], len(hz_arr))
plt.title('Hz')
plt.plot([hz_arr, hz_arr], y1y2)
plt.show()
# bins = hz_arr / high_hz * n_spec
bins = np.floor(hz_arr / high_hz * n_spec)
print('bins', bins.shape, bins.dtype)
filter_banks = np.zeros((n_filters, n_spec), dtype=np.float64)
for m in range(1, n_filters + 1):
    m_minus = int(bins[m - 1])
    m_center = int(bins[m])
    m_plus = int(bins[m + 1])
    for k in range(m_minus, m_center):
        filter_banks[m - 1, k] = (k - bins[m - 1]) / (bins[m] - bins[m - 1])
    for k in range(m_center, m_plus):
        filter_banks[m - 1, k] = (bins[m + 1] - k) / (bins[m + 1] - bins[m])
plt.title('Filters')
for f in filter_banks:
    plt.plot(f)
plt.show()

filtered = np.dot(pow_frames, filter_banks.T)
filtered = np.where(np.isclose(filtered, 0.), np.finfo(float).eps, filtered)
filtered = 20 * np.log10(filtered)
plt.title('Filtered power spectrum')
sns.heatmap(filtered)
plt.show()

num_ceps = 13
mfcc = scipy.fftpack.dct(filtered, type=2, axis=1, norm='ortho')[:, 1: (num_ceps + 1)]
plt.title('DCT')
sns.heatmap(mfcc)
plt.show()
